#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
基于PaddlePaddle的ERNIE-4.5-0.3B-Paddle测试程序
"""

import json
import os
import random
import sys

class PaddleERNIETester:
    def __init__(self, model_path="./ERNIE-4.5-0.3B-Paddle"):
        """
        初始化ERNIE测试器
        
        Args:
            model_path (str): 本地模型路径
        """
        self.model_path = model_path
        self.model = None
        self.tokenizer = None
        self.infinity_instruct_data = None
        
    def check_environment(self):
        """
        检查PaddlePaddle环境
        """
        print("检查PaddlePaddle环境...")
        try:
            import paddle
            print(f"✓ PaddlePaddle版本: {paddle.__version__}")
            return True
        except ImportError:
            print("✗ 未安装PaddlePaddle")
            return False
        except Exception as e:
            print(f"✗ PaddlePaddle环境异常: {e}")
            return False
    
    def load_model_info(self):
        """
        加载模型信息（即使没有PaddlePaddle环境也能显示）
        """
        print("正在加载ERNIE-4.5-0.3B-Paddle模型信息...")
        
        # 检查模型目录是否存在
        if not os.path.exists(self.model_path):
            print(f"✗ 模型目录 {self.model_path} 不存在")
            return False
            
        try:
            # 读取配置文件
            with open(os.path.join(self.model_path, "config.json"), "r", encoding="utf-8") as f:
                config = json.load(f)
                
            print("✓ 模型信息加载成功")
            print(f"  模型类型: {config.get('model_type', '未知')}")
            print(f"  隐藏层大小: {config.get('hidden_size', '未知')}")
            print(f"  层数: {config.get('num_hidden_layers', '未知')}")
            print(f"  注意力头数: {config.get('num_attention_heads', '未知')}")
            
            # 读取tokenizer配置
            with open(os.path.join(self.model_path, "tokenizer_config.json"), "r", encoding="utf-8") as f:
                tokenizer_config = json.load(f)
                
            print(f"  Tokenizer类: {tokenizer_config.get('tokenizer_class', '未知')}")
            
            return True
        except Exception as e:
            print(f"✗ 加载模型信息时出错: {e}")
            return False
    
    def load_infinity_instruct_dataset(self, sample_size=5):
        """
        加载Infinity-Instruct数据集示例
        
        Args:
            sample_size (int): 采样数据集大小
        """
        print("正在加载Infinity-Instruct数据集示例...")
        try:
            # 检查Infinity-Instruct目录
            instruct_dir = "./Infinity-Instruct"
            if not os.path.exists(instruct_dir):
                print("✗ 未找到Infinity-Instruct数据集")
                return False
                
            # 查找json文件
            json_files = [f for f in os.listdir(instruct_dir) if f.endswith('.json')]
            if not json_files:
                print("✗ Infinity-Instruct目录中未找到JSON文件")
                return False
                
            # 读取示例数据
            sample_file = os.path.join(instruct_dir, json_files[0])
            with open(sample_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
                
            # 采样数据
            if len(data) > sample_size:
                self.infinity_instruct_data = random.sample(data, sample_size)
            else:
                self.infinity_instruct_data = data
                
            print(f"✓ 成功加载{len(self.infinity_instruct_data)}条示例数据")
            return True
        except Exception as e:
            print(f"✗ 加载Infinity-Instruct数据集时出错: {e}")
            return False
    
    def simulate_response(self, prompt):
        """
        模拟ERNIE模型的响应
        
        Args:
            prompt (str): 输入提示
            
        Returns:
            str: 模拟的响应
        """
        # 根据输入提示生成相关的模拟响应
        responses = {
            "人工智能": "人工智能是计算机科学的一个分支，它企图了解智能的实质，并生产出一种新的能以人类智能相似的方式做出反应的智能机器。",
            "大语言模型": "大语言模型是一种基于深度学习的自然语言处理模型，能够理解和生成人类语言，广泛应用于文本生成、问答系统等任务。",
            "机器学习": "机器学习是人工智能的一个子领域，通过算法让计算机从数据中学习模式和规律，从而对新数据进行预测或决策。",
            "深度学习": "深度学习是机器学习的一个分支，使用多层神经网络来模拟人脑处理信息的方式，能够自动提取数据的特征。",
            "自然语言处理": "自然语言处理是计算机科学和人工智能领域的重要方向，致力于让计算机理解和生成人类语言。"
        }
        
        # 查找最匹配的关键词
        for key, response in responses.items():
            if key in prompt:
                return response
                
        # 默认响应
        return "这是一个很好的问题。基于我的知识，我可以告诉您这涉及到现代科技的重要领域。随着技术的发展，相关应用正在不断改善我们的生活和工作方式。"
    
    def test_with_instruct_samples(self):
        """
        使用Infinity-Instruct数据集中的样本进行测试
        """
        if not self.infinity_instruct_data:
            print("错误：未加载Infinity-Instruct数据集")
            return
            
        print("\n" + "="*60)
        print("使用Infinity-Instruct数据集样本进行测试")
        print("="*60)
        
        for i, sample in enumerate(self.infinity_instruct_data[:3]):  # 只测试前3个样本
            print(f"\n--- 测试样本 {i+1} ---")
            
            # 获取用户查询
            user_query = None
            for conversation in sample.get("conversations", []):
                if conversation.get("from") == "human":
                    user_query = conversation.get("value")
                    break
                    
            if user_query:
                print(f"输入: {user_query[:100]}..." if len(user_query) > 100 else f"输入: {user_query}")
                
                # 模拟生成响应
                response = self.simulate_response(user_query)
                print(f"模拟响应: {response}")
            else:
                print("未找到有效的用户查询")
    
    def interactive_mode(self):
        """
        交互模式，允许用户输入自定义提示
        """
        print("\n" + "="*60)
        print("交互模式 - 输入 'quit' 退出")
        print("="*60)
        
        while True:
            try:
                user_input = input("\n请输入您的问题: ").strip()
                if user_input.lower() in ['quit', 'exit', '退出']:
                    break
                    
                if user_input:
                    response = self.simulate_response(user_input)
                    print(f"\nERNIE模拟回答:\n{response}")
                else:
                    print("请输入有效问题。")
                    
            except KeyboardInterrupt:
                print("\n\n程序被用户中断。")
                break
            except Exception as e:
                print(f"处理输入时出错: {e}")
    
    def show_installation_guide(self):
        """
        显示安装指南
        """
        print("\n" + "="*60)
        print("PaddlePaddle安装指南")
        print("="*60)
        print("""
要在本地运行ERNIE-4.5-0.3B-Paddle模型，您需要安装PaddlePaddle:

1. 使用pip安装（推荐使用国内镜像源）:
   pip install paddlepaddle -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 如果您使用的是GPU环境:
   pip install paddlepaddle-gpu -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 安装PaddleNLP:
   pip install paddlenlp -i https://pypi.tuna.tsinghua.edu.cn/simple

注意：
- 在Windows环境下安装PaddlePaddle可能会遇到一些问题
- 建议在Linux环境或WSL中安装以获得最佳兼容性
- 如果遇到安装问题，可以考虑使用Docker或在线AI开发平台

安装完成后，您可以使用以下代码加载模型:

import paddle
from paddlenlp.transformers import AutoModelForCausalLM, AutoTokenizer

model_name_or_path = "./ERNIE-4.5-0.3B-Paddle"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True)
        """)
    
    def run_demo(self):
        """
        运行演示程序
        """
        print("ERNIE-4.5-0.3B-Paddle PaddlePaddle测试程序")
        print("="*60)
        
        # 检查环境
        try:
            env_ok = self.check_environment()
        except Exception as e:
            print(f"检查环境时出错: {e}")
            env_ok = False
            
        if not env_ok:
            print("提示：未检测到可用的PaddlePaddle环境")
            self.show_installation_guide()
        
        # 加载模型信息
        self.load_model_info()
        
        # 加载数据集
        self.load_infinity_instruct_dataset()
        
        # 运行预设测试
        self.test_with_instruct_samples()
        
        # 交互模式
        self.interactive_mode()
        
        print("\n感谢使用ERNIE-4.5-0.3B-Paddle测试程序！")

def main():
    """
    主函数
    """
    tester = PaddleERNIETester()
    tester.run_demo()

if __name__ == "__main__":
    main()